package main

import (
	"fmt"

	"github.com/xwb1989/sqlparser"
)

// Parse takes a SQL query string and returns the parsed query statement type
// and table name, or an error if parsing failed.
func Parse(query string) (string, string, error) {
	stmt, err := sqlparser.Parse(query)
	if err != nil {
		return "", "", fmt.Errorf("failed to parse query: %w", err)
	}

	switch stmt := stmt.(type) {
	case *sqlparser.Select:
		return "SELECT", getTableName(stmt.From), nil
	case *sqlparser.Update:
		return "UPDATE", getTableName(stmt.TableExprs), nil
	case *sqlparser.Insert:
		return "INSERT", stmt.Table.Name.String(), nil
	case *sqlparser.Delete:
		return "DELETE", getTableName(stmt.TableExprs), nil
	default:
		return "UNKNOWN", "", fmt.Errorf("unsupported statement type")
	}
}

// getTableName extracts the table name from a SQL node.
func getTableName(node sqlparser.SQLNode) string {
	switch tableExpr := node.(type) {
	case sqlparser.TableName:
		return tableExpr.Name.String()
	case sqlparser.TableExprs:
		for _, expr := range tableExpr {
			if tableName, ok := expr.(*sqlparser.AliasedTableExpr); ok {
				if name, ok := tableName.Expr.(sqlparser.TableName); ok {
					return name.Name.String()
				}
			}
		}
	}
	return ""
}

func main() {
	queries := []string{
		"SELECT * FROM users WHERE id = 1",
		"SELECT id, name FROM users WHERE id = 1",
		"INSERT INTO orders (id, amount) VALUES (1, 100)",
		"UPDATE products SET price = 19.99 WHERE id = 10",
		"DELETE FROM sessions WHERE expired = true",
		"CREATE TABLE logs (id INT, message TEXT)",
	}
	for _, query := range queries {
		fmt.Println("Query: ", query)
		statement, table, err := Parse(query)
		if err != nil {
			fmt.Println("Error:", err)
			continue
		}
		fmt.Printf("Statement: %s, Table: %s\n", statement, table)
	}
}

package main

import (
	"fmt"

	"vitess.io/vitess/go/vt/sqlparser"
)

// Parse takes a SQL query string and returns the parsed query statement type
// and table name(s), or an error if parsing failed.
func Parse(query string) (string, []string, error) {
	p, err := sqlparser.New(sqlparser.Options{})
	if err != nil {
		return "", nil, fmt.Errorf("failed to create parser: %w", err)
	}

	stmt, err := p.Parse(query)
	if err != nil {
		return "", nil, fmt.Errorf("failed to parse query: %w", err)
	}

	var statementType string
	var tables []string

	switch stmt := stmt.(type) {
	case *sqlparser.Select:
		statementType = "SELECT"
		tables = extractTables(stmt.From)
	case *sqlparser.Update:
		statementType = "UPDATE"
		tables = extractTables(stmt.TableExprs)
	case *sqlparser.Insert:
		statementType = "INSERT"
		tables = []string{stmt.Table.TableNameString()}
	case *sqlparser.Delete:
		statementType = "DELETE"
		tables = extractTables(stmt.TableExprs)
	case *sqlparser.CreateTable:
		statementType = "CREATE TABLE"
		tables = []string{stmt.Table.Name.String()}
	case *sqlparser.AlterTable:
		statementType = "ALTER TABLE"
		tables = []string{stmt.Table.Name.String()}
	case *sqlparser.DropTable:
		statementType = "DROP TABLE"
		for _, table := range stmt.FromTables {
			tables = append(tables, table.Name.String())
		}
	case *sqlparser.CreateDatabase:
		statementType = "CREATE DATABASE"
		tables = []string{stmt.DBName.String()}
	case *sqlparser.DropDatabase:
		statementType = "DROP DATABASE"
		tables = []string{stmt.DBName.String()}
	case *sqlparser.TruncateTable:
		statementType = "TRUNCATE TABLE"
		tables = []string{stmt.Table.Name.String()}
	default:
		return "UNKNOWN", nil, fmt.Errorf("unsupported statement type")
	}

	return statementType, tables, nil
}

// extractTables extracts table names from a list of SQL nodes.
func extractTables(exprs sqlparser.TableExprs) []string {
	var tables []string
	for _, expr := range exprs {
		switch tableExpr := expr.(type) {
		case *sqlparser.AliasedTableExpr:
			if name, ok := tableExpr.Expr.(sqlparser.TableName); ok {
				tables = append(tables, name.Name.String())
			}
		}
	}
	return tables
}

func main() {
	queries := []string{
		"SELECT * FROM users WHERE id = 1",
		"SELECT id, name FROM users WHERE id = 1",
		"INSERT INTO users (id, name) VALUES (1, 'Alice')",
		"UPDATE users SET name = 'Bob' WHERE id = 1",
		"DELETE FROM users WHERE id = 1",
		"CREATE TABLE users (id INT, name VARCHAR(100))",
		"ALTER TABLE users ADD COLUMN age INT",
		"DROP TABLE users",
		"CREATE DATABASE test_db",
		"DROP DATABASE test_db",
		"TRUNCATE TABLE users",
	}

	for _, query := range queries {
		fmt.Println("Query: ", query)
		statement, tables, err := Parse(query)
		if err != nil {
			fmt.Printf("Error parsing query: %s\nQuery: %s\n\n", err, query)
			continue
		}
		fmt.Printf("Statement: %s, Tables/DBs: %v\n", statement, tables)
	}
}